import torch
from torchvision.transforms import transforms
import h5py
from PIL import Image
import time

def predictImages(image, torch_weight):

    # preprocess
    transform_test = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.40272543], std=[0.13901867])
        ])

    DEVICE = torch.device("cpu")

    image = Image.open(image).convert("RGB")
    image = transform_test(image)
    image = torch.unsqueeze(image, dim=0).to(DEVICE)


    # Predict
    model = torch.load(torch_weight)
    model = model.eval().to(DEVICE)
    start_time = time.time()
    for i in range(300):
        out = model(image)
    end_time = time.time()
    print((end_time - start_time) / 300)
    return

if __name__ == '__main__':
    image = 'img.png'
    torch_weight = 'convnext_s.pth'
    predictImages(image, torch_weight)

    '''
    pth cpu
    convnext_s 224 0.12801373839378358
    swinV2_s 256 0.33226001024246216
   
    '''